package main

import (
	"sync/atomic"
)

type ITask interface {
	run()
	post()
}

type TaskPool struct {
	InputChan      chan ITask
	OutputChan     chan ITask
	RunningThreads int32
}

func CreateTaskPool(threads int32, inputBuf int, outputBuf int) *TaskPool {
	pool := &TaskPool{}

	pool.InputChan = make(chan ITask, inputBuf)
	pool.OutputChan = make(chan ITask, outputBuf)
	pool.RunningThreads = threads

	for i := int32(0); i < threads; i++ {
		go pool.runner()
	}

	return pool
}

func (self *TaskPool) runner() {
	for task := range self.InputChan {
		task.run()
		self.OutputChan <- task
	}

	if atomic.AddInt32(&self.RunningThreads, -1) == 0 {
		close(self.OutputChan)
	}
}

func (self *TaskPool) CloseInput() {
	close(self.InputChan)
}

func (self *TaskPool) Join(onFinish func()) {
	for result := range self.OutputChan {
		result.post()
	}

	if onFinish != nil {
		onFinish()
	}
}
